# -*- coding: utf-8 -*-

import torchvision.transforms as tv_transforms


def build_transforms(cfg):
    phase = cfg.PHASE

    if phase == "train":
        transforms = tv_transforms.Compose([
            tv_transforms.RandomCrop(size=cfg.TRANSFORM.CROP_SIZE,
                                     padding=cfg.TRANSFORM.CROP_PADDING_SIZE),
            tv_transforms.RandomHorizontalFlip(),
            tv_transforms.ToTensor(),
            tv_transforms.Normalize(mean=cfg.INPUT.PIXEL_MEAN,
                                    std=cfg.INPUT.PIXEL_STD)
        ])
    else:
        transforms = tv_transforms.Compose([
            tv_transforms.ToTensor(),
            tv_transforms.Normalize(mean=cfg.INPUT.PIXEL_MEAN,
                                    std=cfg.INPUT.PIXEL_STD)
        ])
    return transforms
